#pragma once

#include <algorithm/analytical_model/algorithm/cuda/inter_op/Data.h>
#include <utils/Concepts.h>
#include <utils/config.h>

#include <tuple>
#include <type_traits>

// Provide CPU versions ?

namespace solar
{
    namespace cuda
    {
        void sumFluxMap(const solar::cuda::inter_op::Int2& resolution, float*& result,
                        FUNC_CONST_INPUT float* input, cudaStream_t stream = nullptr,
                        bool delete_input = false);

        void divFluxMap(const solar::cuda::inter_op::Int2& resolution, float*& result,
                        const float& num, cudaStream_t stream = nullptr);
    } // namespace cuda

    /// @brief 累加Flux分布
    /// @param resolution 分辨率
    /// @param result 累加结果（on device）
    /// @param input 需要累加的Flux分布（on device）
    template <ConstSpanLike Container>
    void sumFluxMaps(const solar::cuda::inter_op::Int2& resolution, float*& result,
                     FUNC_CONST_INPUT Container& input, cudaStream_t stream = nullptr)
    {
        static_assert(std::is_same_v<float*, std::remove_reference_t<decltype(*input.data())>>,
                      "Input should be a container of float* (on device)");
        for (auto& current_input : input)
        {
            cuda::sumFluxMap(resolution, result, current_input, stream);
        }
    }

    /// @brief 计算Flux分布的平均值
    /// @param resolution 分辨率
    /// @param result 平均值结果（on device）
    /// @param input 需要计算平均值的Flux分布（on device）
    template <ConstSpanLike Container>
    void avgFluxMaps(const solar::cuda::inter_op::Int2& resolution, float*& result,
                     FUNC_CONST_INPUT Container& input, const float& num, cudaStream_t stream = nullptr)
    {
        sumFluxMaps(resolution, result, input, stream);
        cuda::divFluxMap(resolution, result, num, stream);
    }

    /// @brief 查找FluxMap上的最大值及其坐标
    /// @param resolution 分辨率
    /// @param input FluxMap数据（on device）
    /// @return (最大值(单位：W), 最大值下标（在input中）)
    auto findFluxMapMax(const solar::cuda::inter_op::Int2& resolution,
                        FUNC_CONST_INPUT float* input) -> std::tuple<float, std::uint64_t>;

    constexpr auto
    convertToFluxMapIndex(const solar::cuda::inter_op::Int2& resolution,
                          const std::uint64_t& index_1d) -> solar::cuda::inter_op::Int2
    {
        return {static_cast<int>(resolution.y - 1 - index_1d / resolution.x),
                static_cast<int>(index_1d % resolution.x)};
    }

    /// @brief 计算FluxMap上的能量之和
    /// @param resolution 分辨率
    /// @param input FluxMap数据（on device）
    auto totalEnergyOnFluxMap(const solar::cuda::inter_op::Int2& resolution,
                              FUNC_CONST_INPUT float* input) -> float;
} // namespace solar
